#include "ulib.h"

struct va_item_t
{
	 unsigned int flag : 2;
	 unsigned int idx  : 30;
};

#define VA_FLAG_FH   0
#define VA_FLAG_FB   1
#define VA_FLAG_FT   2
#define VA_FLAG_USED 3

struct va_ctrl_t
{
	 spin_lock_t lock;
	 uint32_t    free_head;
};

static char __va_item_size_assert[sizeof(struct va_item_t) == 4 ? 0 : -1];

extern volatile struct va_ctrl_t va_ctrl;
extern volatile struct va_item_t va_items[PAGETAB_COUNT * PAGEDIR_COUNT];

int
va_init(void)
{
	 int result;
	 if ((result = sc_alloc_temp_4M((void *)UHHEAD)) < 0)
		  return result;

	 spl_init(&va_ctrl.lock);
	 va_ctrl.free_head = PAGE_NUM(UHEAP);

	 int head = PAGE_NUM(UHEAP);
	 int tail = PAGE_NUM(UWLIM) - 1;
	 
	 va_items[head].flag = VA_FLAG_FH;
	 va_items[head].idx	 = 0;
	 va_items[tail].flag = VA_FLAG_FT;
	 va_items[tail].idx  = 0;

	 va_items[head + 1].flag = VA_FLAG_FB;
	 va_items[head + 1].idx  = tail - 1;
	 va_items[tail - 1].flag = VA_FLAG_FB;
	 va_items[tail - 1].idx  = head + 1;

	 return result;
}

void *
va_alloc(int size)
{
	 // sc_kprintf("va_alloc start size = %d\n", size);
	 
	 if (size & (PAGE_SIZE - 1))
		  size = (size >> PAGE_SHIFT) + 1;
	 else size = (size >> PAGE_SHIFT);
	 /* minimum size of 2 * PAGE_SIZE */
	 if (size < 2) size = 2;
	 
	 int result, last_tail, tail;
	 sc_spl_acquire(&va_ctrl.lock);
	 last_tail = 0;
	 result = va_ctrl.free_head;
	 while (result != 0)
	 {
		  /* Get tail of free block */
		  tail = va_items[result + 1].flag == VA_FLAG_FB ?
			   va_items[result + 1].idx + 1 : result + 1;

		  // sc_kprintf("Looking at %d - %d\n", result, tail);
		  
		  int _size = tail - result + 1;
		  /* Do not alloc the 1 page size fragment */
		  if (_size == size || _size == size + 1)
		  {
			   /* Remove the free block and catch the result */
			   if (last_tail == 0)
					va_ctrl.free_head = va_items[tail].idx;
			   else va_items[last_tail].idx = va_items[tail].idx;

			   if (va_items[tail].idx != 0)
					va_items[va_items[tail].idx].idx = 0;

			   break;
		  }
		  else if (_size > size)
		  {
			   /* Split the free block then break */
			   /* The new tail */
			   int n_tail = tail - size;
			   if (n_tail - 1 != result)
			   {
					va_items[n_tail - 1].flag = VA_FLAG_FB;
					va_items[n_tail - 1].idx  = result + 1;
			   }
			   va_items[n_tail] = va_items[tail];
			   if (result + 1 < n_tail - 1)
					va_items[result + 1].idx = n_tail - 1;
			   result = n_tail + 1;

			   break;
		  }
		  else
		  {
			   /* Not enough space, try next one on the list */
			   result = va_items[tail].idx;
			   last_tail = tail;
		  }
	 }
	 spl_release(&va_ctrl.lock);

	 // sc_kprintf("va_alloc ends with result idx = %d\n", result);
	 if (result != 0)
	 {
		  va_items[result].flag = VA_FLAG_USED;
		  va_items[result].idx  = tail;
		  va_items[tail].flag   = VA_FLAG_USED;
		  va_items[tail].idx    = result;

		  return (void *)(result << PAGE_SHIFT);
	 }
	 else return NULL;
}

void
va_free(void *addr)
{
	 int head = PAGE_NUM(addr);
	 int tail = va_items[head].idx;

	 // sc_kprintf("va_free %d - %d\n", head, tail);
	 
	 sc_spl_acquire(&va_ctrl.lock);
	 /* Merge up */
	 if (head > PAGE_NUM(UHEAP) &&
		 va_items[head - 1].flag != VA_FLAG_USED)
	 {
		  va_items[tail] = va_items[head - 1];
		  head = va_items[head - 2].flag == VA_FLAG_FB ?
			   va_items[head - 2].idx - 1 : head - 2;

		  // sc_kprintf("merge up! head = %d,\n", head);
	 }
	 else
	 {
		  va_items[head].flag = VA_FLAG_FH;
		  va_items[head].idx  = 0;
		  
		  va_items[tail].flag = VA_FLAG_FT;
		  va_items[tail].idx  = 0;
	 }

	 /* Merge down */
	 if (tail + 1 < PAGE_NUM(UWLIM) &&
		 va_items[tail + 1].flag != VA_FLAG_USED)
	 {
		  int n_tail = va_items[tail + 2].flag == VA_FLAG_FB ?
			   va_items[tail + 2].idx + 1: tail + 2;

		  /* remove block from list */
		  if (va_items[tail + 1].idx == 0)
			   va_ctrl.free_head = va_items[n_tail].idx;
		  else va_items[va_items[tail + 1].idx].idx = va_items[n_tail].idx;
		  if (va_items[n_tail].idx != 0)
			   va_items[va_items[n_tail].idx].idx = va_items[tail + 1].idx;

		  va_items[n_tail] = va_items[tail];
		  tail = n_tail;

		  // sc_kprintf("merge down! tail = %d,\n", tail);
	 }

	 if (head == PAGE_NUM(addr))
	 {
		  /* insert into list */
		  va_items[head].idx = 0;
		  if ((va_items[tail].idx = va_ctrl.free_head) != 0)
			   va_items[va_ctrl.free_head].idx = head;
		  va_ctrl.free_head = head;
	 }

	 /* Set the inside block info */
	 if (head + 1 != tail)
	 {
		  va_items[head + 1].flag = VA_FLAG_FB;
		  va_items[head + 1].idx  = tail - 1;
		  va_items[tail - 1].flag = VA_FLAG_FB;
		  va_items[tail - 1].idx  = head - 1;
	 }
	 spl_release(&va_ctrl.lock);

	 // sc_kprintf("va_free finished\n");
}

void *bmalloc(int size)
{
	 char *result = va_alloc(size);
	 if (result != NULL)
	 {
		  char *cur;
		  for (cur = result; cur < result + size; cur += PAGE_SIZE)
			   sc_alloc_temp_4K(cur);
	 }
	 return result;
}

void bfree(void *addr)
{
	 va_free(addr);
}
